
#include "select_v2_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"

constexpr uint32_t BLOCK_SIZE = 32;
constexpr int32_t inputVarNum = 3;
constexpr int32_t maxDimNum = 64;

namespace optiling {
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{

    SelectV2TilingData tiling;
    //const gert::StorageShape* x1_shape = context->GetInputShape(1);
    uint32_t totallength = 1;
    // for (int i = 0; i < x1_shape->GetStorageShape().GetDimNum(); i++)
    //     totallength *= x1_shape->GetStorageShape().GetDim(i);
    // tiling.set_totallength(totallength);

    uint64_t ubSize;
    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);
    uint32_t inputNum = context->GetInputShape(0)->GetStorageShape().GetShapeSize(); 
    uint32_t inputBytes = GetSizeByDataType(context->GetInputDesc(1)->GetDataType()); 
    uint32_t inputLength = inputBytes * inputNum;
    std::cout << "inputBytes:" << inputBytes << std::endl;
    uint32_t ubDataNumber = 16;
    if(inputBytes == 2) {
        ubDataNumber = 9;
    } else if(inputBytes == 4) {
        ubDataNumber = 9;
    } else if(inputBytes == 1) {
        ubDataNumber = 20;
    }
    uint32_t tileBlockNum = (ubSize / BLOCK_SIZE) / ubDataNumber; //每个ub段可用的空间块数
    uint32_t tileDataNum = (tileBlockNum * BLOCK_SIZE) / inputBytes; //每次处理的数据量
    tileDataNum = tileDataNum / BLOCK_SIZE * BLOCK_SIZE;
    tiling.set_tileDataNum(tileDataNum);

    uint32_t condition_length = context->GetInputShape(0)->GetStorageShape().GetShapeSize();
    uint32_t x1_length = context->GetInputShape(1)->GetStorageShape().GetShapeSize();
    uint32_t x2_length = context->GetInputShape(2)->GetStorageShape().GetShapeSize();
    int64_t shape[maxDimNum * inputVarNum], shapefull[maxDimNum];
    bool e1flag = false;
    int64_t numshapes = 1;
    int64_t numshapes_condition = context->GetInputShape(0)->GetStorageShape().GetDimNum();
    int64_t numshapes1 = context->GetInputShape(1)->GetStorageShape().GetDimNum();
    int64_t numshapes2 = context->GetInputShape(2)->GetStorageShape().GetDimNum();
    std::cout << "numshapes_condition:" << numshapes_condition << std::endl;
    std::cout << "numshapes1:" << numshapes1 << std::endl;
    std::cout << "numshapes2:" << numshapes2 << std::endl;
    numshapes = (numshapes1 > numshapes2) ? numshapes1 : numshapes2;
    numshapes = (numshapes > numshapes_condition) ? numshapes : numshapes_condition;
    for (int k = 0; k < maxDimNum * inputVarNum; ++k) {
        shape[k] = 1;
    }    
    for (int k = 0; k < maxDimNum; ++k) {
        shapefull[k] = 1;
    }    
    for (int k = 0; k < inputVarNum; ++k) {
        int64_t *ss = &shape[k * maxDimNum];
        const gert::StorageShape* inputshape = context->GetInputShape(k);
        for (int i = 0; i < inputshape->GetStorageShape().GetDimNum(); i++) {
            ss[i] = inputshape->GetStorageShape().GetDim(i);
        }
    }
    for (int k = 0; k < numshapes; ++k) {
        int64_t *ss = &shape[0];
        int64_t *sf = &shapefull[0];
        sf[k] = (ss[k] > ss[k + maxDimNum]) ? ss[k] : ss[k + maxDimNum];  
        sf[k] = (sf[k] > ss[k + (maxDimNum * 2)]) ? sf[k] : ss[k + (maxDimNum * 2)];  
        totallength *=  sf[k];
    }

    if(condition_length == 1) {
        context->SetTilingKey(100);
        const gert::Tensor *condition_tensor = context->GetInputTensor(0);
        if (condition_tensor == nullptr) {
            return ge::GRAPH_FAILED;
        }
        auto addrC = condition_tensor->GetData<bool>();
        e1flag = *addrC;
    } else if (condition_length < totallength || x1_length < totallength || x2_length < totallength) {
        context->SetTilingKey(200);
    } else {
        context->SetTilingKey(300);
    }
    tiling.set_totallength(totallength);
    tiling.set_e1flag(e1flag);
    tiling.set_condition_length(condition_length);
    tiling.set_x1_length(x1_length);
    tiling.set_x2_length(x2_length);
    tiling.set_numshapes(numshapes);
    tiling.set_shape(shape);
    tiling.set_shapefull(shapefull);

    std::cout << "tiling_key:" << context->GetTilingKey() << std::endl;
    std::cout << "totallength:" << totallength << std::endl;
    std::cout << "tileDataNum:" << tileDataNum << std::endl;
    std::cout << "e1flag:" << e1flag << std::endl;
    std::cout << "condition_length:" << condition_length << std::endl;
    std::cout << "x1_length:" << x1_length << std::endl;
    std::cout << "x2_length:" << x2_length << std::endl;
    std::cout << "numshapes:" << numshapes << std::endl;

    context->SetBlockDim(1);
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());

    return ge::GRAPH_SUCCESS;
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
}


namespace ops {
class SelectV2 : public OpDef {
public:
    explicit SelectV2(const char* name) : OpDef(name)
    {
        this->Input("condition")
            .ParamType(REQUIRED)
            .DataType({ge::DT_BOOL, ge::DT_BOOL, ge::DT_BOOL, ge::DT_BOOL})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("x1")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT8, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("x2")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT8, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT8, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});

        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend310b");

    }
};

OP_ADD(SelectV2);
}
